from pgmpy.models import DiscreteBayesianNetwork  # 使用正确的类名
from pgmpy.factors.discrete import TabularCPD
from pgmpy.inference import VariableElimination

# 创建贝叶斯网络结构
model = DiscreteBayesianNetwork([
    ('季节', '感冒'),     # 季节 → 感冒
    ('感冒', '发烧'),     # 感冒 → 发烧  
    ('感冒', '流鼻涕')    # 感冒 → 流鼻涕
])

# 定义条件概率分布
# 季节的概率分布（先验概率）
cpd_季节 = TabularCPD(
    variable='季节', 
    variable_card=4,  # 4个季节
    values=[[0.25], [0.25], [0.25], [0.25]]  # 每个季节概率相等
)

# 给定季节条件下感冒的概率
cpd_感冒 = TabularCPD(
    variable='感冒',
    variable_card=2,  # 0:不感冒, 1:感冒
    values=[
        [0.9, 0.7, 0.6, 0.8],   # 不感冒的概率（春,夏,秋,冬）
        [0.1, 0.3, 0.4, 0.2]    # 感冒的概率
    ],
    evidence=['季节'],
    evidence_card=[4]  # 季节有4种可能取值[2](@ref)[5](@ref)
)

# 给定感冒条件下发烧的概率
cpd_发烧 = TabularCPD(
    variable='发烧',
    variable_card=2,  # 0:不发烧, 1:发烧
    values=[
        [0.95, 0.3],   # 不感冒时不发烧概率95%，感冒时30%
        [0.05, 0.7]    # 不感冒时发烧概率5%，感冒时70%
    ],
    evidence=['感冒'],
    evidence_card=[2]  # 感冒有2种可能取值（是/否）[2](@ref)[5](@ref)
)

# 给定感冒条件下流鼻涕的概率
cpd_流鼻涕 = TabularCPD(
    variable='流鼻涕',
    variable_card=2,  # 0:不流鼻涕, 1:流鼻涕
    values=[
        [0.9, 0.2],    # 不感冒时不流鼻涕概率90%，感冒时20%
        [0.1, 0.8]     # 不感冒时流鼻涕概率10%，感冒时80%
    ],
    evidence=['感冒'],
    evidence_card=[2]  # 感冒有2种可能取值[2](@ref)[5](@ref)
)

# 将概率分布添加到模型中
model.add_cpds(cpd_季节, cpd_感冒, cpd_发烧, cpd_流鼻涕)

# 检查模型是否配置正确
print("模型检查:", model.check_model())

# 进行概率推理
inference = VariableElimination(model)

print("=== 医疗诊断贝叶斯网络推理结果 ===\n")

# 例子1：已知是冬季，计算感冒的概率
result1 = inference.query(variables=['感冒'], evidence={'季节': 3})  # 3代表冬季
print("1. 冬季感冒概率分析:")
print(f"   不感冒概率: {result1.values[0]:.2%}")
print(f"   感冒概率: {result1.values[1]:.2%}")

# 例子2：已知患者流鼻涕，推断感冒的概率
result2 = inference.query(variables=['感冒'], evidence={'流鼻涕': 1})
print("\n2. 流鼻涕时感冒概率分析:")
print(f"   不感冒概率: {result2.values[0]:.2%}")
print(f"   感冒概率: {result2.values[1]:.2%}")

# 例子3：已知发烧且流鼻涕，推断感冒的概率
result3 = inference.query(variables=['感冒'], evidence={'发烧': 1, '流鼻涕': 1})
print("\n3. 发烧且流鼻涕时感冒概率分析:")
print(f"   不感冒概率: {result3.values[0]:.2%}")
print(f"   感冒概率: {result3.values[1]:.2%}")

# 额外例子：不同季节的感冒概率对比
print("\n4. 不同季节感冒概率对比:")
for season in range(4):
    season_names = ['春', '夏', '秋', '冬']
    result = inference.query(variables=['感冒'], evidence={'季节': season})
    print(f"   {season_names[season]}季感冒概率: {result.values[1]:.2%}")
